import taichi as ti

# 初始化太极
ti.init(arch=ti.gpu)

# ODOP
@ti.data_oriented
class SolarSystem:
    def __init__(self,n,dt):
        self.n = n
        self.dt = dt
        self.x = ti.Vector.field(2,dtype=ti.f32,shape=n)
        self.v = ti.Vector.field(2,dtype=ti.f32,shape=n)
        # 恒星（中心）位置
        self.center = ti.Vector.field(2,dtype=ti.f32,shape=())

    @staticmethod
    @ti.func
    def random_around(center,radius):
        # center +- r
        return center+radius*(ti.random()-0.5)*2
    
    @ti.kernel
    def initialize(self):
        for i in range(self.n):
            # 随机生成偏移 y方向：0.15-0.45
            offset = ti.Vector([0.0,self.random_around(0.3,0.15)])
            # 行星位置，恒星正上方
            self.x[i] = self.center[None] + offset
            # -x方向的初始速度，切线方向
            self.v[i] = [-offset[1],offset[0]]
            self.v[i] *= 1.5/offset.norm()
    
    @ti.func
    def gravity(self,pos):
        # 万有引力加速度（忽略G，M）
        offset = -(pos-self.center[None])
        return offset / offset.norm()**3
    
    @ti.kernel
    def integrate(self):
        # 前向欧拉
        for i in range(self.n):
            # v[t+1] = v[t] + dt* (r*GM/r3)
            self.v[i] += self.dt * self.gravity(self.x[i])
            # x[t+1] = x[t] + dt*v[t+1]
            self.x[i] += self.dt * self.v[i]

# 9个行星，dt=0.0005
solar = SolarSystem(9,0.0005)
# 恒星位置（屏幕中心）
solar.center[None] = [0.5,0.5]
# 初始化行星参数
solar.initialize()

gui = ti.GUI("solar system",background_color=0x25A6D9)

while True:
    if gui.get_event():
        if gui.event.key == gui.SPACE and gui.event.type == gui.PRESS:
            solar.initialize()
    for i in range(10):
        solar.integrate()
    gui.circle(solar.center.to_numpy(),radius=20,color=0x8C274C)
    gui.circles(solar.x.to_numpy(),radius=5,color=0xFFFFFF)
    gui.show()
